package com.example.qaplatform.service;

import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import java.util.*;

@Service
public class RagflowService {
    @Value("${ragflow.api.key}")
    private String ragflowApiKey;

    @Value("${ragflow.dataset.id}")
    private String datasetId;

    // 获取图片URL（不变）
    public String getImageUrl(String name) {
        try {
            RestTemplate restTemplate = new RestTemplate();
            String docListUrl = "http://localhost/api/v1/datasets/" + datasetId + "/documents?page=1&page_size=10&keywords=" + name;
            HttpHeaders docHeaders = new HttpHeaders();
            docHeaders.set("Authorization", "Bearer " + ragflowApiKey);
            HttpEntity<String> docEntity = new HttpEntity<>(docHeaders);
            ResponseEntity<Map> docListResp = restTemplate.exchange(docListUrl, HttpMethod.GET, docEntity, Map.class);
            if (docListResp.getStatusCode() == HttpStatus.OK && docListResp.getBody() != null) {
                Object dataObj = docListResp.getBody().get("data");
                if (dataObj instanceof Map) {
                    Object docsObj = ((Map) dataObj).get("docs");
                    if (docsObj instanceof List && !((List) docsObj).isEmpty()) {
                        System.out.println("[getImageUrl] docsObj: " + docsObj);
                        for (Object docObj : (List) docsObj) {
                            if (docObj instanceof Map) {
                                String docName = ((Map) docObj).getOrDefault("name", "").toString();
                                String docId = ((Map) docObj).getOrDefault("id", "").toString();
                                System.out.println("[getImageUrl] docName: " + docName + ", docId: " + docId);
                                // 只查图片文档
                                if (docName.endsWith(".jpg") || docName.endsWith(".png")) {
                                    String chunkListUrl = "http://localhost/api/v1/datasets/" + datasetId + "/documents/" + docId + "/chunks?page=1&page_size=10";
                                    HttpHeaders chunkHeaders = new HttpHeaders();
                                    chunkHeaders.set("Authorization", "Bearer " + ragflowApiKey);
                                    HttpEntity<String> chunkEntity = new HttpEntity<>(chunkHeaders);
                                    ResponseEntity<Map> chunkListResp = restTemplate.exchange(chunkListUrl, HttpMethod.GET, chunkEntity, Map.class);
                                    if (chunkListResp.getStatusCode() == HttpStatus.OK && chunkListResp.getBody() != null) {
                                        Object chunkDataObj = chunkListResp.getBody().get("data");
                                        if (chunkDataObj instanceof Map) {
                                            Object chunksObj = ((Map) chunkDataObj).get("chunks");
                                            if (chunksObj instanceof List) {
                                                for (Object chunkObj : (List) chunksObj) {
                                                    if (chunkObj instanceof Map) {
                                                        System.out.println("[getImageUrl] chunkObj: " + chunkObj);
                                                        String imageId = ((Map) chunkObj).getOrDefault("image_id", "").toString();
                                                        if (!imageId.isEmpty()) {
                                                            return "http://localhost/v1/document/image/" + imageId;
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    // 获取百科简介（用关键词查知识库，不查图片docId）
    public String getDesc(String name) {
        try {
            RestTemplate restTemplate = new RestTemplate();
            String url = "http://localhost/api/v1/chats_openai/abbf981662dc11f0b68a0242ac120005/chat/completions";
            Map<String, Object> body = new HashMap<>();
            body.put("model", "qwen2.5vl:3b");
            List<Map<String, String>> messages = new ArrayList<>();
            Map<String, String> userMsg = new HashMap<>();
            userMsg.put("role", "user");
            userMsg.put("content", name);
            messages.add(userMsg);
            body.put("messages", messages);
            body.put("stream", false);

            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("Authorization", "Bearer " + ragflowApiKey);
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(body, headers);

            ResponseEntity<Map> resp = restTemplate.postForEntity(url, entity, Map.class);
            if (resp.getStatusCode() == HttpStatus.OK && resp.getBody() != null) {
                List choices = (List) resp.getBody().get("choices");
                if (choices != null && !choices.isEmpty()) {
                    Map choice = (Map) choices.get(0);
                    Map message = (Map) choice.get("message");
                    String content = message.get("content").toString();
                    // 去除 [ID:xxx] 这类内容
                    content = content.replaceAll("\\[ID:[^\\]]*\\]", "");
                    return content.trim();
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }
} 